import os
import numpy as np
import matplotlib.pyplot as plt
from tensorflow.keras.datasets import mnist

def extract_mnist_to_disk(output_dir='mnist_data'):
    """
    将MNIST数据集从mnist.npz解压到本地磁盘
    
    参数:
        output_dir: 输出目录路径(默认'mnist_data')
    """
    # 创建输出目录(如果不存在)
    os.makedirs(output_dir, exist_ok=True)
    
    # 加载MNIST数据集(会自动下载或从缓存读取)
    (x_train, y_train), (x_test, y_test) = mnist.load_data()
    
    # 保存为.npy文件
    np.save(os.path.join(output_dir, 'x_train.npy'), x_train)
    np.save(os.path.join(output_dir, 'y_train.npy'), y_train)
    np.save(os.path.join(output_dir, 'x_test.npy'), x_test)
    np.save(os.path.join(output_dir, 'y_test.npy'), y_test)
    
    # 创建可视化样本目录
    samples_dir = os.path.join(output_dir, 'samples')
    os.makedirs(samples_dir, exist_ok=True)
    
    # 保存前10个训练样本为图片
    for i in range(10):
        plt.imshow(x_train[i], cmap='gray')
        plt.title(f"Label: {y_train[i]}")
        plt.axis('off')
        plt.savefig(os.path.join(samples_dir, f'train_sample_{i}_label_{y_train[i]}.png'))
        plt.close()
    
    print(f"MNIST数据已成功导出到: {os.path.abspath(output_dir)}")
    print(f"包含: 4个.npy文件和10个样本图片")

if __name__ == '__main__':
    # 使用默认输出目录'mnist_data'
    extract_mnist_to_disk()
    
    # 或者指定自定义目录
    # extract_mnist_to_disk('/path/to/your/directory')